package com.sharding.demo.config;

import com.alibaba.ttl.TransmittableThreadLocal;
import com.google.common.base.Preconditions;
import com.google.common.collect.HashMultimap;
import com.google.common.collect.Multimap;
import lombok.AccessLevel;
import lombok.NoArgsConstructor;
import org.apache.shardingsphere.infra.hint.HintManager;

import java.util.Collection;
import java.util.Collections;

/**
 * @Author hu xu hong
 * @Date 2023/2/1 15:16
 */
@NoArgsConstructor(access = AccessLevel.PRIVATE)
public class HintManagerTTL implements AutoCloseable{

    private static final ThreadLocal<HintManagerTTL> HINT_MANAGER_HOLDER = new TransmittableThreadLocal<>();

    private final Multimap<String, Comparable<?>> databaseShardingValues = HashMultimap.create();

    private final Multimap<String, Comparable<?>> tableShardingValues = HashMultimap.create();

    private boolean databaseShardingOnly;

    private boolean masterRouteOnly;

    /**
     * Get a new instance for {@code HintManager}.
     *
     * @return  {@code HintManager} instance
     */
    public static HintManagerTTL getInstance() {
        Preconditions.checkState(null == HINT_MANAGER_HOLDER.get(), "Hint has previous value, please clear first.");
        HintManagerTTL result = new HintManagerTTL();
        HINT_MANAGER_HOLDER.set(result);
        return result;
    }

    /**
     * Set sharding value for database sharding only.
     *
     * <p>The sharding operator is {@code =}</p>
     *
     * @param value sharding value
     */
    public void setDatabaseShardingValue(final Comparable<?> value) {
        databaseShardingValues.clear();
        tableShardingValues.clear();
        databaseShardingValues.put("", value);
        databaseShardingOnly = true;
    }

    /**
     * Add sharding value for database.
     *
     * <p>The sharding operator is {@code =}</p>
     *
     * @param logicTable logic table name
     * @param value sharding value
     */
    public void addDatabaseShardingValue(final String logicTable, final Comparable<?> value) {
        if (databaseShardingOnly) {
            databaseShardingValues.removeAll("");
        }
        databaseShardingValues.put(logicTable, value);
        databaseShardingOnly = false;
    }

    /**
     * Add sharding value for table.
     *
     * <p>The sharding operator is {@code =}</p>
     *
     * @param logicTable logic table name
     * @param value sharding value
     */
    public void addTableShardingValue(final String logicTable, final Comparable<?> value) {
        if (databaseShardingOnly) {
            databaseShardingValues.removeAll("");
        }
        tableShardingValues.put(logicTable, value);
        databaseShardingOnly = false;
    }

    /**
     * Get database sharding values.
     *
     * @return database sharding values
     */
    public static Collection<Comparable<?>> getDatabaseShardingValues() {
        return getDatabaseShardingValues("");
    }

    /**
     * Get database sharding values.
     *
     * @param logicTable logic table
     * @return database sharding values
     */
    public static Collection<Comparable<?>> getDatabaseShardingValues(final String logicTable) {
        return null == HINT_MANAGER_HOLDER.get() ? Collections.emptyList() : HINT_MANAGER_HOLDER.get().databaseShardingValues.get(logicTable);
    }

    /**
     * Get table sharding values.
     *
     * @param logicTable logic table name
     * @return table sharding values
     */
    public static Collection<Comparable<?>> getTableShardingValues(final String logicTable) {
        return null == HINT_MANAGER_HOLDER.get() ? Collections.emptyList() : HINT_MANAGER_HOLDER.get().tableShardingValues.get(logicTable);
    }

    /**
     * Judge whether database sharding only.
     *
     * @return database sharding or not
     */
    public static boolean isDatabaseShardingOnly() {
        return null != HINT_MANAGER_HOLDER.get() && HINT_MANAGER_HOLDER.get().databaseShardingOnly;
    }

    /**
     * Set database operation force route to master database only.
     */
    public void setMasterRouteOnly() {
        masterRouteOnly = true;
    }

    /**
     * Judge whether route to master database only or not.
     *
     * @return route to master database only or not
     */
    public static boolean isMasterRouteOnly() {
        return null != HINT_MANAGER_HOLDER.get() && HINT_MANAGER_HOLDER.get().masterRouteOnly;
    }

    /**
     * Clear threadlocal for hint manager.
     */
    public static void clear() {
        HINT_MANAGER_HOLDER.remove();
    }

    @Override
    public void close() {
        HintManager.clear();
    }
}
